package com.zzs.judge.rabbit;

import com.zzs.common.core.constant.RabbitConsts;
import com.zzs.common.core.entity.msg.CodeSubmitMessage;
import com.zzs.judge.config.JavaContainerPool;
import com.zzs.judge.service.JudgeCodeService;
import jakarta.annotation.PostConstruct;
import org.springframework.amqp.rabbit.annotation.Exchange;
import org.springframework.amqp.rabbit.annotation.Queue;
import org.springframework.amqp.rabbit.annotation.QueueBinding;
import org.springframework.amqp.rabbit.annotation.RabbitListener;
import org.springframework.amqp.rabbit.listener.RabbitListenerEndpointRegistry;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;


/**
 * @author zzs
 * @date 2025/04/01
 */
@Component
public class MQListener {
    @Autowired
    private JudgeCodeService judgeCodeService;

    @Autowired
    private JavaContainerPool javaContainerPool;

    private ThreadPoolExecutor threadPoolExecutor;

    @Autowired
    private RabbitListenerEndpointRegistry registry;


    @PostConstruct
    private void init() {
//        核心线程数 = 最大线程数 = 容器池大小
        int corePoolSize = javaContainerPool.getContainerSize();
        long keepAliveTime = 60L;
        int queueCapacity = 100;
        threadPoolExecutor = new ThreadPoolExecutor(
                corePoolSize,
                corePoolSize,
                keepAliveTime,
                TimeUnit.SECONDS,
                new ArrayBlockingQueue<>(queueCapacity),
                new ThreadPoolExecutor.AbortPolicy()
        );
    }


    @RabbitListener(bindings = @QueueBinding(
            value = @Queue(name = RabbitConsts.JAVA_QUEUE),
            exchange = @Exchange(name = RabbitConsts.JAVA_EXCHANGE),
            key = {RabbitConsts.JAVA_ROUTING_KEY}),
            id = "javaQueueListener"
    )
    public void javaQueueListener(CodeSubmitMessage codeSubmitMessage) {
        if (isThreadPoolBusy()) {
//            阻塞队列快满了, 暂停rabbitmq的消费
            pauseConsuming();
            return;
        }
        threadPoolExecutor.execute(() -> {
            String containerId = null;
            try {
                containerId = javaContainerPool.getContainerId();
                judgeCodeService.judgeJavaCode(containerId, codeSubmitMessage);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            } finally {
                if (containerId != null) {
                    javaContainerPool.releaseContainerId(containerId);
                }
                if (isThreadPoolFree()) {
//                    如果阻塞队列空闲了, 恢复rabbitmq的消费
                    resumeConsuming();
                }
            }
        });
    }

    private boolean isThreadPoolBusy() {
//        此时阻塞队列的元素个数
        int queueSize = threadPoolExecutor.getQueue().size();
        int queueCapacity = threadPoolExecutor.getQueue().remainingCapacity() + queueSize;
        return queueSize >= queueCapacity * 0.8;
    }

    private boolean isThreadPoolFree() {
//        此时阻塞队列的元素个数
        int queueSize = threadPoolExecutor.getQueue().size();
        int queueCapacity = threadPoolExecutor.getQueue().remainingCapacity() + queueSize;
        return queueSize <= queueCapacity * 0.2;
    }

    /**
     * 暂停从 RabbitMQ 消费消息
     */
    public void pauseConsuming() {
        registry.getListenerContainer("javaQueueListener").stop();
    }

    /**
     * 恢复从 RabbitMQ 消费消息
     */
    public void resumeConsuming() {
        registry.getListenerContainer("javaQueueListener").start();
    }
}
